"""
@author : linrh
@homepage : https://gitee.com/linrh-DUT
@version: 1.0.0
@when : 2023/5/17
@file: test.py
"""
import torch

from conf import *
from data import *
from models.model.transformer import Transformer

x = loader._tensor_transform('I love you', loader.tokenize_source, loader.vocab_source).unsqueeze(dim=0).to(device)
print(x)
model = Transformer(src_pad_idx=src_pad_idx,
                    trg_pad_idx=trg_pad_idx,
                    trg_sos_idx=trg_sos_idx,
                    d_model=d_model,
                    enc_voc_size=enc_voc_size,
                    dec_voc_size=dec_voc_size,
                    max_len=max_len,
                    ffn_hidden=ffn_hidden,
                    n_head=n_heads,
                    n_layers=n_layers,
                    drop_prob=drop_prob,
                    device=device).to(device)

def load_weight(model):
    model.load_state_dict(torch.load("./saved/model-saved.pt"))

load_weight(model=model)
s = [loader.BOS_IDX]

for _ in range(20):
    y = torch.Tensor([s]).long().to(device)
    output = model(x, y)
    output_reshape = output.contiguous().view(-1, output.shape[-1])
    y = output_reshape.argmax(dim=-1)
    s = [loader.BOS_IDX] + y.tolist()
    print(s)
    print(loader.get_sentence(s, loader.vocab_target))
    if s[-1] == 3:
        break
